// 构建向量数据库，可以用faiss-node

import { config } from "dotenv";
import { TextLoader } from "langchain/document_loaders/fs/text";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { useScoreThresholdRetriever ,useMultiQueryRetriever} from "./utils";
import { OpenAIEmbeddings } from "@langchain/openai";
import { FaissStore } from "@langchain/community/vectorstores/faiss";
import path from "path";
import { fileURLToPath } from "url";

// ES Module 中获取 __dirname 的方法
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

await config();
const runner = async () => {
  const loader = new TextLoader("data/kong.txt");
  const docs = await loader.load();
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 100,
    chunkOverlap: 20,
  });

  const splitDocs = await splitter.splitDocuments(docs);
  console.log(splitDocs[0]);

  const embeddings = new OpenAIEmbeddings({
    apiKey: process.env.DASHSCOPE_API_KEY,
    model: "text-embedding-v1", // 阿里云的 embedding 模型
    batchSize: 25, // 阿里云限制每次最多处理 25 个文档
    configuration: {
      baseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
    },
  });
  const vectorStore = await FaissStore.fromDocuments(splitDocs, embeddings);
  const dir = path.resolve(__dirname, "..", "..", "data", "db", "kongyiji");
  await vectorStore.save(dir);
};

// runner();

const loader = async () => {
  const dir = path.resolve(__dirname, "..", "..", "data", "db", "kongyiji");
  const embeddings = new OpenAIEmbeddings({
    apiKey: process.env.DASHSCOPE_API_KEY,
    model: "text-embedding-v1", // 阿里云的 embedding 模型
    batchSize: 25, // 阿里云限制每次最多处理 25 个文档
    configuration: {
      baseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
    },
  });
  const vectorStore = await FaissStore.load(dir, embeddings);
  // 我们可以看到这里，asRetriever（3）这里的这个 3，其实是很难抉择的，是否可以优化呢，可以的，ScoreThresholdRetriever
    //   const retriever = vectorStore.asRetriever(3);

  // 这下面的这个方法就是封装了一下 ScoreThresholdRetriever，方便我们使用，传入 vectorStore 即可
  const retriever = useScoreThresholdRetriever(vectorStore);
  const multiQueryRetriever = useMultiQueryRetriever(vectorStore);
  const res = await retriever.invoke("所以，到底偷没偷东西呀");
  const multiQueryRes = await multiQueryRetriever.invoke("所以，到底偷没偷东西呀");
  console.log(res);
  console.log(multiQueryRes);
};



loader();
